package com.yifeng.repo.storage.elasticsearch;

import com.google.common.base.Strings;
import com.google.common.collect.Maps;
import com.yifeng.repo.base.consts.DeployModeEnum;
import com.yifeng.repo.storage.elasticsearch.configure.ElasticsearchProperties;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.nio.conn.ssl.SSLIOSessionStrategy;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.net.ssl.*;
import java.security.SecureRandom;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * Created by daibing on 2021/5/17.
 */
@Slf4j
public class ElasticSearchClient {
    private final ElasticsearchProperties properties;

    private final Map<String, RestHighLevelClient> clientMap = Maps.newConcurrentMap();

    public ElasticSearchClient(ElasticsearchProperties properties) {
        this.properties = properties;
    }

    public void init() {
        // 兼容老版本单个
        if (CollectionUtils.isEmpty(properties.getClusters())){
            RestClientBuilder restClientBuilder = initClient(properties);
            String aliasName = StringUtils.hasText(properties.getAliasName()) ? properties.getAliasName() : "default";
            clientMap.put(aliasName, new RestHighLevelClient(restClientBuilder));
        }
        // 识别多集群
        else {
            for (ElasticsearchProperties cluster : properties.getClusters()) {
                RestClientBuilder restClientBuilder = initClient(cluster);
                String aliasName = StringUtils.hasText(cluster.getAliasName()) ? cluster.getAliasName() : "default";
                clientMap.put(aliasName, new RestHighLevelClient(restClientBuilder));
            }
        }
    }

    private RestClientBuilder initClient(ElasticsearchProperties properties) {
        List<HttpHost> httpHosts = this.buildHttpHost(properties);
        return RestClient.builder(httpHosts.toArray(new HttpHost[]{}))
                .setRequestConfigCallback(requestConfigBuilder -> {
                    requestConfigBuilder.setConnectTimeout(properties.getConnectTimeout());
                    requestConfigBuilder.setSocketTimeout(properties.getSocketTimeout());
                    requestConfigBuilder.setConnectionRequestTimeout(properties.getConnectionRequestTimeout());
                    return requestConfigBuilder;
                }).setHttpClientConfigCallback(httpClientBuilder -> {
                    // httpClientBuilder.disableAuthCaching();
                    httpClientBuilder.setKeepAliveStrategy((response, context) -> properties.getKeepAliveTimeout());
                    httpClientBuilder.setMaxConnTotal(properties.getMaxConnTotal());
                    httpClientBuilder.setMaxConnPerRoute(properties.getMaxConnPerRoute());
                    if (!Strings.isNullOrEmpty(properties.getUserName()) && !Strings.isNullOrEmpty(properties.getPassword())) {
                        BasicCredentialsProvider provider = new BasicCredentialsProvider();
                        provider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(properties.getUserName(), properties.getPassword()));
                        httpClientBuilder.setDefaultCredentialsProvider(provider);
                    }
                    if (properties.getScheme().equalsIgnoreCase("https")) {
                        httpClientBuilder.setSSLStrategy(buildSSLStrategy());
                    }
                    return httpClientBuilder;
                });
    }

    public void destroy() {
        try {
            for (RestHighLevelClient client : clientMap.values()) {
                client.close();
            }
        } catch (Throwable t) {
            log.warn("destroy failed: ", t);
        }
    }

    public RestHighLevelClient getClient() {
        String aliasName = StringUtils.hasText(properties.getAliasName()) ? properties.getAliasName() : "default";
        return this.clientMap.get(aliasName);
    }

    public RestHighLevelClient getClient(String aliasName) {
        return this.clientMap.get(aliasName);
    }

    public ElasticsearchProperties getProperties() {
        return properties;
    }

    private List<HttpHost> buildHttpHost(ElasticsearchProperties properties) {
        if (properties.getDeployMode() == DeployModeEnum.STANDALONE) {
            return Collections.singletonList(new HttpHost(properties.getHostname(), properties.getPort(), properties.getScheme()));
        }
        if (properties.getDeployMode() == DeployModeEnum.CLUSTER) {
            String[] hostAndPorts = properties.getClusterHostAndPorts().split(",");
            List<HttpHost> httpHostList = new ArrayList<>(hostAndPorts.length);
            for (String hostAndPort : hostAndPorts) {
                String[] split = hostAndPort.split(":");
                HttpHost httpHost = new HttpHost(split[0], Integer.parseInt(split[1]), properties.getScheme());
                httpHostList.add(httpHost);
            }
            return httpHostList;
        }
        throw new RuntimeException("elasticsearch deploy mode config failed! " + properties.getDeployMode());
    }

    private SSLIOSessionStrategy buildSSLStrategy() {
        TrustManager[] trustAllCerts = new TrustManager[]{
                new X509TrustManager() {
                    @Override
                    public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {
                    }

                    @Override
                    public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException {
                    }

                    @Override
                    public X509Certificate[] getAcceptedIssuers() {
                        return null;
                    }
                }
        };
        HostnameVerifier trustAnyHostnameVerifier = new HostnameVerifier() {
            @Override
            public boolean verify(String s, SSLSession sslSession) {
                return true;
            }
        };
        try {
            SSLContext sc = SSLContext.getInstance("SSL");
            sc.init(null, trustAllCerts, new SecureRandom());
            return new SSLIOSessionStrategy(sc, trustAnyHostnameVerifier);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

}
